Improved and simplified BinaryOperation with "stubborn" location inference #1599

glwagner · 2021-04-20T03:45:46Z

This PR changes the structure of BinaryOperations and the way that locations are inferred when constructing them.

Before this PR, BinaryOperations stored three interpolation operators: two for each argument, and a third interpolation function applied after the binary operation was computed. The motivation for this design was to ensure that products of fields were always computed at their shared location, while allowing abstract operations to be built at user-specified locations. In other words, users might want something like

using Oceananigans
grid = RegularRectilinearGrid(size=(1, 1, 1), extent=(1, 1, 1))
model = IncompressibleModel(architecture=CPU(), grid=grid)
u, v, w = model.velocities

uu = @at (Center, Center, Center) u * u

# output
BinaryOperation at (Center, Center, Center)
├── grid: RegularRectilinearGrid{Float64, Periodic, Periodic, Bounded}(Nx=1, Ny=1, Nz=1)
│   └── domain: x ∈ [0.0, 1.0], y ∈ [0.0, 1.0], z ∈ [-1.0, 0.0]
└── tree: 
    * at (Center, Center, Center) via ℑxᶜᵃᵃ
    ├── Field located at (Face, Center, Center)
    └── Field located at (Face, Center, Center)

which would compute u * u at (Face, Center, Center), and subsequently interpolate to cell centers. The object uu would then be defined at cell centers.

The main issue of this design in terms of functionality is that it produces expression trees that interpolate "too eagerly". A common example is a turbulent kinetic energy computation:

U = AveragedField(u, dims=(1, 2))
V = AveragedField(v, dims=(1, 2))

tke = @at (Center, Center, Center) ((u - U)^2 + (v - V)^2 + w^2)

Inspection of this object reveals that while u - U does not interpolate, the squaring (u - U)^2 would be performed at cell centers:

julia> tke.args[1].b
2

julia> tke.args[1].a
BinaryOperation at (Center, Center, Center)
├── grid: RegularRectilinearGrid{Float64, Periodic, Periodic, Bounded}(Nx=1, Ny=1, Nz=1)
│   └── domain: x ∈ [0.0, 1.0], y ∈ [0.0, 1.0], z ∈ [-1.0, 0.0]
└── tree: 
    - at (Center, Center, Center) via identity
    ├── Field located at (Face, Center, Center)
    └── AveragedField over dims=(1, 2) located at (⋅, ⋅, Center) of Field located at (Face, Center, Center)

This isn't what we want (usually): instead, we want both u - U and subsequent squaring in (u - U)^2 to be performed at (Face, Center, Center).

This PR addresses this issue by getting rid of the third interpolation in BinaryOperation, and changing how locations for BinaryOperation are inferred. Now, when locations are specified via @at, they are taken as a "suggestion" that only acts if the two elements of the binary operation have different concrete locations such that a difference needs to be resolved. In other cases (such as a binary operation between fields at common locations, or a binary operation between a field and a number), the location of the members of the operation is preserved. In this way, BinaryOperations are "stubborn".

We thus have results like

julia> *((Center, Center, Center), u, u)
BinaryOperation at (Face, Center, Center)
├── grid: RegularRectilinearGrid{Float64, Periodic, Periodic, Bounded}(Nx=1, Ny=1, Nz=1)
│   └── domain: x ∈ [0.0, 1.0], y ∈ [0.0, 1.0], z ∈ [-1.0, 0.0]
└── tree: 
    * at (Face, Center, Center)
    ├── Field located at (Face, Center, Center)
    └── Field located at (Face, Center, Center)

and

julia> tke = @at (Center, Center, Center) ((u - U)^2 + (v - V)^2 + w^2);

julia> tke.arg.args[1]
BinaryOperation at (Face, Center, Center)
├── grid: RegularRectilinearGrid{Float64, Periodic, Periodic, Bounded}(Nx=1, Ny=1, Nz=1)
│   └── domain: x ∈ [0.0, 1.0], y ∈ [0.0, 1.0], z ∈ [-1.0, 0.0]
└── tree: 
    ^ at (Face, Center, Center)
    ├── - at (Face, Center, Center)
    │   ├── Field located at (Face, Center, Center)
    │   └── AveragedField over dims=(1, 2) located at (⋅, ⋅, Center) of Field located at (Face, Center, Center)
    └── 2

An added benefit is that the BinaryOperation object is simpler. This could help compilation as well...

Of course, users do want to be able to specify the location of BinaryOperations for output. For this we change how @at works: now we wrap the entire user-prescribed expression in an interpolation operator that interpolates the result to the user-specified location. If the location of the underlying expression is already at the user-specified location, this is just an identity. But when the underlying operation is a "stubborn" BinaryOperator, we interpolate:

julia> uu = @at (Center, Center, Center) u * u
UnaryOperation at (Center, Center, Center)
├── grid: RegularRectilinearGrid{Float64, Periodic, Periodic, Bounded}(Nx=1, Ny=1, Nz=1)
│   └── domain: x ∈ [0.0, 1.0], y ∈ [0.0, 1.0], z ∈ [-1.0, 0.0]
└── tree: 
    identity at (Center, Center, Center) via ℑxᶜᵃᵃ
    └── * at (Face, Center, Center)
        ├── Field located at (Face, Center, Center)
        └── Field located at (Face, Center, Center)

julia> uu.arg
BinaryOperation at (Face, Center, Center)
├── grid: RegularRectilinearGrid{Float64, Periodic, Periodic, Bounded}(Nx=1, Ny=1, Nz=1)
│   └── domain: x ∈ [0.0, 1.0], y ∈ [0.0, 1.0], z ∈ [-1.0, 0.0]
└── tree: 
    * at (Face, Center, Center)
    ├── Field located at (Face, Center, Center)
    └── Field located at (Face, Center, Center)

This also means that operations between ReducedField and Field also end up at the right place, and we have no need to throw an error if a BinaryOperation has a Nothing location (in fact, this might be a useful abstraction for binary operations between ReducedField).

This reverts commit 1a77150.

ali-ramadhan

🤞 this helps with compilation! I guess this is technically a breaking change? Sounds like it might be fine for LESbrary.jl though.

glwagner · 2021-04-20T18:58:33Z

It is very much breaking! And for the better. I'll bump version. But yeah, since it mostly changes complex abstract operations that previously haven't compiled on the GPU, most code isn't affected.

tomchor · 2021-04-20T20:45:02Z

Sorry for not following, but I don't see what breaks here for the end user! Could someone give a quick example/explanation?

glwagner · 2021-04-21T00:54:07Z

Sorry for not following, but I don't see what breaks here for the end user! Could someone give a quick example/explanation?

Apologies that my explanation was not clear. It's a breaking change because the same user input, such as

tke = @at (Center, Center, Center) ((u - U)^2 + (v - V)^2 + w^2) / 2

produces a different object after this update:

Before this PR, u - U would be interpolated to (Center, Center, Center), and then the binary operation ^(2, u - U) would be calculated at (Center, Center, Center).
After this PR, both u - U and ^(2, u - U) = (u - U)^2 are calculated at (Face, Center, Center). Interpolation is then performed to (Center, Center, Center) after the exponentiation to form the three-part sum.

Let me know if that makes sense or if another example would be useful.

tomchor · 2021-04-21T01:01:58Z

@glwagner Thanks! I actually had gotten that, but I was under the impression that the previous syntax wasn't gonna work. (Which would possibly break many scripts I have...)

Thanks for the explanation though!

(Btw, this PR was super timely. I was planning on opening an issue this week about the order of operations in abstract operations following our discussion out it the past week 👍)

glwagner · 2021-04-21T01:05:58Z

@glwagner Thanks! I actually had gotten that, but I was under the impression that the previous syntax wasn't gonna work. (Which would possibly break many scripts I have...)

Thanks for the explanation though!

(Btw, this PR was super timely. I was planning on opening an issue this week about the order of operations in abstract operations following our discussion out it the past week 👍)

This was a hard problem to solve. We'll see if this algorithm is sufficient for all cases. Other solutions are welcome.

glwagner · 2021-04-21T01:11:26Z

under the impression that the previous syntax wasn't gonna work

I believe this qualifies as a "breaking change" because any code that relied on consistent output from an abstract operation could, in theory, break (for example, a regression test that passes only when output remains constant). It does not change the API however, just the results.

glwagner · 2021-04-21T01:14:19Z

After some tweaks it does indeed look like

tke_ccc  = @at (Center, Center, Center) ((u - U)^2  + (v - V)^2 + w^2) / 2

compiles on the GPU! What a journey.

Release notes: * Tests and fixes for `FFTBasedPoissonSolver` for topologies with `Flat` dimensions (#1560) * Improved `AbstractOperations` that are much more likely to compile on the GPU, with better "location inference" for `BinaryOperation` (#1595, #1599)

glwagner added 10 commits April 17, 2021 15:43

More general location inference for binary operations

fc50632

Add identity to list of unary operators

d5c944b

Adds additional type parameter to zero field

1a77150

Revert "Adds additional type parameter to zero field"

5d2d5ef

This reverts commit 1a77150.

Introduces "stubborn" binary operations

1380894

Cleans up doc strings and adapt_structure for AbstractOperations

2abcae5

Fix some scoping issues with interpolation wrapper for at macro

b68c486

Move interpolate_operation definition to at.jl

92785d8

No need to throw errors for operations at Nothing location

d407611

Computations with AveragedField might work?

8d81ef1

ali-ramadhan approved these changes Apr 20, 2021

View reviewed changes

glwagner added 3 commits April 20, 2021 12:05

Adds incremental build up of AveragedField operations testing

b767230

Dont interpolate unecessarily

4da9637

Bugfix in abstract operations tests

9c04e5a

Use custom binary op function for getindex

ca5717e

No more test skip in abstrsct operations tests!

4a0f36f

glwagner merged commit 91e40bf into master Apr 21, 2021

glwagner deleted the glw/less-invasive-at branch April 21, 2021 03:37

glwagner mentioned this pull request Apr 21, 2021

Bump to 0.55.0 #1600

Merged

glwagner mentioned this pull request Apr 23, 2021

Added information about CuArrays #1609

Merged

glwagner mentioned this pull request Jul 22, 2021

Complex AbstractOperations cannot be computed on GPU #1241

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Improved and simplified BinaryOperation with "stubborn" location inference #1599

Improved and simplified BinaryOperation with "stubborn" location inference #1599

glwagner commented Apr 20, 2021 •

edited

Loading

ali-ramadhan left a comment

glwagner commented Apr 20, 2021

tomchor commented Apr 20, 2021

glwagner commented Apr 21, 2021 •

edited

Loading

tomchor commented Apr 21, 2021

glwagner commented Apr 21, 2021

glwagner commented Apr 21, 2021

glwagner commented Apr 21, 2021

Improved and simplified BinaryOperation with "stubborn" location inference #1599

Improved and simplified BinaryOperation with "stubborn" location inference #1599

Conversation

glwagner commented Apr 20, 2021 • edited Loading

ali-ramadhan left a comment

Choose a reason for hiding this comment

glwagner commented Apr 20, 2021

tomchor commented Apr 20, 2021

glwagner commented Apr 21, 2021 • edited Loading

tomchor commented Apr 21, 2021

glwagner commented Apr 21, 2021

glwagner commented Apr 21, 2021

glwagner commented Apr 21, 2021

glwagner commented Apr 20, 2021 •

edited

Loading

glwagner commented Apr 21, 2021 •

edited

Loading